#include <string.h>
#include <utility>
#include "regwrap.h"
#include <shlwapi.h>

struct keypair
{
    const char *szName;
    HKEY hRootKey;
};

static keypair rkam[] =
{
    { "HKLM",                  HKEY_LOCAL_MACHINE },
    { "HKCU",                  HKEY_CURRENT_USER },
    { "HKCR",                  HKEY_CLASSES_ROOT },
    { "HKEY_LOCAL_MACHINE",    HKEY_LOCAL_MACHINE },
    { "HKEY_CURRENT_USER",     HKEY_CURRENT_USER },
    { "HKEY_CLASSES_ROOT",     HKEY_CLASSES_ROOT },
    { "HKEY_USERS",            HKEY_USERS },
    { "HKEY_CURRENT_CONFIG",   HKEY_CURRENT_CONFIG },
    { "HKEY_PERFORMANCE_DATA", HKEY_PERFORMANCE_DATA },
    { "HKEY_DYN_DATA",         HKEY_DYN_DATA },
    { 0, 0 }
};

void RegWrap::Init( HKEY hRootKey, const tstring& sSubKey, bool bCreate )
{
    hkSession = 0;

    // Open/create the registry key

    nInitResult = bCreate ? ::RegCreateKeyEx( hRootKey, sSubKey.c_str(), 0, "", 0, KEY_ALL_ACCESS, 0, &hkSession, 0 )
                          : ::RegOpenKeyEx( hRootKey, sSubKey.c_str(), 0, KEY_READ, &hkSession );

    if ( nInitResult != ERROR_SUCCESS )
        hkSession = 0;
}

RegWrap::RegWrap( const tstring& sKey, bool bCreate ) :
    hkSession( 0 )
{
    // Parse the key

    unsigned long iFirstSlash = sKey.find_first_of( "\\" );

    if ( iFirstSlash == tstring::npos )
        return;

    tstring sRootKey = sKey.substr( 0, iFirstSlash );
    tstring sSubKey = sKey.substr( iFirstSlash+1 );

    // Get the root key by its name

    HKEY hRootKey = (HKEY)-1;

    for ( unsigned long i = 0; rkam[i].szName != 0; ++i )
        if ( sRootKey == rkam[i].szName )
            hRootKey = rkam[i].hRootKey;

    if ( hRootKey == (HKEY)-1 )
        return;

    Init( hRootKey, sSubKey, bCreate );
}

RegWrap::~RegWrap()
{
    if ( !this->operator!() )
        ::RegCloseKey( hkSession );
}

tstring RegWrap::ReadString( const tstring& sName, const char *szDefaultValue ) const
{
    unsigned long dwValueType;
    unsigned char buf[4096];
    unsigned long dwBufSize = sizeof buf - 1;

    // Read the value into buffer

    if ( ReadToBuf(sName, dwValueType, buf, dwBufSize) != 0 || dwValueType != REG_SZ )
        return szDefaultValue ? szDefaultValue : "";

    return (char*)buf;
}

unsigned long RegWrap::ReadDword( const tstring& sName, unsigned long dwDefaultValue ) const
{
    unsigned long dwValueType;
    unsigned long dwValue;
    unsigned long dwBufSize = sizeof(dwValue);

    // Read the value into buffer

    if ( ReadToBuf(sName, dwValueType, (unsigned char*)&dwValue, dwBufSize) != 0 || dwValueType != REG_DWORD )
        return dwDefaultValue;

    return dwValue;
}

long RegWrap::ReadToBuf( const tstring& sName, unsigned long &dwValueType, unsigned char *buf, unsigned long &dwBufSize ) const
{
    if ( this->operator!() || buf == 0 )
        return -1;

    // Read the value

    long nResult = ::RegQueryValueEx( hkSession, sName.c_str(), 0, &dwValueType, buf, &dwBufSize );

    if ( nResult != ERROR_SUCCESS )
        return -1;

    if ( dwValueType == REG_SZ || dwValueType == REG_MULTI_SZ || dwValueType == REG_EXPAND_SZ )
        buf[dwBufSize] = '\0';

    return 0;
}


HRESULT RegWrap::WriteString( const tstring& sName, const tstring& sValue ) const
{
    if ( this->operator!() )
        return HRESULT_FROM_WIN32( nInitResult );

    long nResult = ::RegSetValueEx( hkSession, sName.c_str(), 0, REG_SZ, (const BYTE *)sValue.c_str(), sValue.length() );
    return HRESULT_FROM_WIN32( nResult );
}

HRESULT RegWrap::WriteDword( const tstring& sName, unsigned long dwValue ) const
{
    if ( this->operator!() )
        return HRESULT_FROM_WIN32( nInitResult );

    long nResult = ::RegSetValueEx( hkSession, sName.c_str(), 0, REG_DWORD, (const BYTE *)&dwValue, sizeof(dwValue) );
    return HRESULT_FROM_WIN32( nResult );
}

HRESULT RegWrap::DeleteValue( const tstring& sName ) const
{
    if ( this->operator!() )
        return HRESULT_FROM_WIN32( nInitResult );

    long nResult = ::RegDeleteValue( hkSession, sName.c_str() );
    return HRESULT_FROM_WIN32( nResult );
}

HRESULT RegWrap::DeleteKey( const tstring& sName ) const
{
    if ( this->operator!() )
        return HRESULT_FROM_WIN32( nInitResult );

    long nResult = ::SHDeleteKey( hkSession, sName.c_str() );
    return HRESULT_FROM_WIN32( nResult );
}

HRESULT RegWrap::EnumKeys( std::vector<tstring>& v ) const
{
    v.clear();

    if ( this->operator!() )
        return HRESULT_FROM_WIN32( nInitResult );

    long nResult = 0;

    unsigned char szBuf[4096];

    for ( unsigned long i = 0; ; ++i )
    {
        unsigned long dwBufSize = sizeof szBuf;

        nResult = ::RegEnumKeyEx( hkSession, i, (LPTSTR)szBuf, &dwBufSize, 0, 0, 0, 0 );

        if ( nResult == ERROR_SUCCESS || nResult == ERROR_MORE_DATA ) {
            v.push_back( (char*)szBuf );
            continue;
        }

        if ( nResult == ERROR_NO_MORE_ITEMS ) {
            nResult = S_OK;
            break;
        }
    }
    
    return HRESULT_FROM_WIN32( nResult );
}

HRESULT RegWrap::EnumValues( std::vector<tstring>& v ) const
{
    v.clear();

    if ( this->operator!() )
        return HRESULT_FROM_WIN32( nInitResult );

    long nResult = 0;

    unsigned char szBuf[4096];

    for ( unsigned long i = 0; ; ++i )
    {
        unsigned long dwBufSize = sizeof szBuf;

        nResult = ::RegEnumValue( hkSession, i, (LPTSTR)szBuf, &dwBufSize, 0, 0, 0, 0 );

        if ( nResult == ERROR_SUCCESS || nResult == ERROR_MORE_DATA ) {
            v.push_back( (char*)szBuf );
            continue;
        }

        if ( nResult == ERROR_NO_MORE_ITEMS ) {
            nResult = S_OK;
            break;
        }
    }

    return HRESULT_FROM_WIN32( nResult );
}
